import os
import cv2
import matplotlib.pyplot as plt
from skimage import data, draw, color, transform, feature, measure
import SimpleITK as sitk
import numpy as np
import time
import math

from . import geoutils, fileutils

__all__ = ['find_long_axis', 'get_exist_range', 'draw_axis']

'''
issues:
    1. label 可能会不稳定，即横断位连续时，矢状位不一定连续。对心尖、平均切线有影响
    2. 调查几何计算的库？
    3. 输入分割不轴对齐时，实际轴向并非切线的角分线，故会不准确。考虑：
        1）用包围盒预处理
        2）迭代旋转求长轴
'''

global_i = 0


def get_tangent(gray_image) -> np.ndarray:
    """
    输入灰度图，基于 Sobel算子计算梯度，进而得到切向量图并返回。对于图像梯度为0的位置，其切向量无意义，规定为[1, 0]

    :param gray_image: 待计算的灰度图
    :return (gray_image.shape, 2) 的 ndarray
    """
    temp_image = gray_image.astype(dtype=np.float32)    # 转换数据格式

    # Scharr算子的核心比较小（3X3），非常局部时比较准，但画出的切线在整体看来不太像，大概率穿过心脏
    # dx_image = cv2.Scharr(temp_image, cv2.CV_32F, 1, 0, borderType=cv2.BORDER_REPLICATE)
    # dy_image = cv2.Scharr(temp_image, cv2.CV_32F, 0, 1, borderType=cv2.BORDER_REPLICATE)

    # 使用 7X7 的Sobel核心
    dx_image = cv2.Sobel(temp_image, cv2.CV_32F, 1, 0, ksize=7, borderType=cv2.BORDER_REPLICATE)
    dy_image = cv2.Sobel(temp_image, cv2.CV_32F, 0, 1, ksize=7, borderType=cv2.BORDER_REPLICATE)

    tan_image = []
    for row in range(gray_image.shape[0]):
        tan_image.append([])
        for col in range(gray_image.shape[1]):
            # grad(p) = （dp/dx, dp/dy）, (x, y) 顺序下
            # 切向量(tangent) dot 梯度(即法向量) = 0
            tangent = np.array([-dx_image[row, col], dy_image[row, col]])
            # 单位化
            # 梯度为 0 的位置，规定切向量为（1, 0），且使用不到（因为只使用边界的切向量）
            if np.linalg.norm(tangent) == 0:
                tan_image[-1].append([1, 0])
            else:
                tangent /= np.linalg.norm(tangent)
                tan_image[-1].append(tangent)
    ret = np.asarray(tan_image)
    return ret


def draw_tangent(ax, point, tangent, half_length):
    """
    在图 ax 中画一条切线（给定点、切向、长度）

    :param ax: ax to draw on
    :param point: middle point of the line, the order is (y, x), aka (row, col)
    :param tangent: tangent vector to determine the direction of the line
    :param half_length: half of the line's length
    :return: None
    """
    point0 = point + half_length * tangent
    point1 = point - half_length * tangent
    ax.add_line(plt.Line2D([point0[1], point1[1]], [point0[0], point1[0]]))
    ax.plot([point[1]], [point[0]], 'o')


def get_epical(tan_image, edge_vertices, ax=None, epsilon=0.05, delta=5) -> list:
    """
    找到心尖范围，在图 ax 中画出

    :param tan_image: the tangent image
    :param edge_vertices: the edge pixels in the tan_image
    :param ax: the ax to draw epical on. If None, then no drawing
    :param epsilon: the precise degree for epical
    :param delta: the y-distance of the epical range
    :return: a list of epical pixels
    """
    candidates = []
    horizontal_vec = np.asarray([0, 1])
    R = np.asarray([
        [math.cos(math.pi / 12), -math.sin(math.pi / 12)],
        [math.sin(math.pi / 12), math.cos(math.pi / 12)]
    ])
    horizontal_vec = np.matmul(R, horizontal_vec)
    # print(horizontal_vec)
    for i, j in edge_vertices:
        if 1 - abs(np.vdot(tan_image[i, j], horizontal_vec)) < epsilon: # 切向与水平向量几乎平行时，加入待选点集
            candidates.append([i, j])

    candidates.sort(key=lambda x: x[0]) # 按 y 的值（即行数）排序
    tail = 0
    # warning: must have at least 2 candidates, or the loop will not be entered
    for i in range(1, len(candidates)):
        if candidates[i][0] - candidates[i-1][0] > delta:   # 若相邻点的垂直距离大于 delta，则跳出，寻找完毕
            break
        tail = i + 1
    # print('count of apical vertices:', tail)

    if ax is not None:
        ax.scatter([x[1] for x in candidates[:tail]], [x[0] for x in candidates[:tail]], s=10, c=[[0, 1, 0]], marker='.')

    return candidates[:tail]


def get_overall_tangent(tan_image: np.ndarray, outer_left, outer_right, ratio_start=0.25, ratio_end=0.75, delta_x=30, ax=None) -> list:
    """
    以左下点的垂直高度为底，最高点为顶，取（ratio_start ~ ratio_end）的区间，作为平均切线的计算区间

    :param tan_image: tangent image
    :param outer_left: left edge pixels, the order must be from bottom to top
    :param outer_right: right edge pixels, the order must be from bottom to top
    :param ratio_start: start ratio in [lb, top)
    :param ratio_end: end ratio in [lb, top)
    :param ax: the ax to draw onto
    :return: two (mean_point, mean_tan) for left and right relatively
    """

    global global_i

    left_bottom_index = 0
    # 若无突变，则左下角视为整体的最下位置
    for i in range(1, round(len(outer_left) * ratio_end)):
        if outer_left[i - 1][1] - outer_left[i][1] > delta_x:   # 左边界水平位置突变处，认为是左下点
            left_bottom_index = i
            break

    # 可视化选取的计算区间
    # total_height = outer_left[left_bottom_index][0] - outer_left[-1][0]
    # overall_range_y = [outer_left[-1][0] + 0.25 * total_height, outer_left[left_bottom_index][0] - 0.25 * total_height]
    # ax.hlines(overall_range_y, [0, 0], [512, 512])

    total_index_len = len(outer_left) - left_bottom_index
    # [ , )，以索引表示的计算区间
    overall_range_index = [left_bottom_index + ratio_start * total_index_len, left_bottom_index + ratio_end * total_index_len]

    # print("left range is ", (outer_left[overall_range_index[0]], outer_left[overall_range_index[1] - 1]))
    # print("right range is ", (outer_right[overall_range_index[0]], outer_right[overall_range_index[1] - 1]))

    # 平均的计算过程
    mean_tan_left = np.zeros(2)
    mean_tan_right = np.zeros(2)
    mean_point_left = np.zeros(2)
    mean_point_right = np.zeros(2)
    count_left = 0
    count_right = 0

    # 跳过不连续的 outer
    last_left = outer_left[round(overall_range_index[0])]
    last_right = outer_right[round(overall_range_index[0])]
    epsilon_left = 5
    epsilon_right = 5

    # 跳过不光滑的 outer
    last_tan_left = np.zeros(2)
    last_tan_right = np.zeros(2)
    cur_tan_left = np.zeros(2)
    cur_tan_right = np.zeros(2)
    cur_point_left = np.zeros(2)
    cur_point_right = np.zeros(2)
    strip_left = 0
    strip_right = 0
    strip_max = 10
    delta_theta = 0.25 * math.pi

    for i in range(round(overall_range_index[0]), round(overall_range_index[1])):
        left_point = outer_left[i]
        right_point = outer_right[i]

        if abs(left_point[1] - last_left[1]) < epsilon_left:
            # mean_tan_left += tan_image[left_point[0], left_point[1]]
            # mean_point_left += left_point
            cur_tan_left += tan_image[left_point[0], left_point[1]]
            cur_point_left += left_point

            # count_left += 1
            strip_left += 1

            last_left = left_point
            epsilon_left = 5
        epsilon_left += 1

        if abs(right_point[1] - last_right[1]) < epsilon_right:
            # mean_tan_right += tan_image[right_point[0], right_point[1]]
            # mean_point_right += right_point
            cur_tan_right += tan_image[right_point[0], right_point[1]]
            cur_point_right += right_point

            # count_right += 1
            strip_right += 1

            last_right = right_point
            epsilon_right = 5
        epsilon_right += 1

        if i == round(overall_range_index[1]) - 1 or strip_left == strip_max:
            if (last_tan_left != np.zeros(2)).any():
                try:
                    if np.dot(last_tan_left, cur_tan_left) / (np.linalg.norm(last_tan_left) * np.linalg.norm(cur_tan_left))\
                            > math.cos(delta_theta):
                        mean_tan_left += cur_tan_left
                        mean_point_left += cur_point_left
                        count_left += strip_left
                        last_tan_left = cur_tan_left
                        cur_tan_left = np.zeros(2)
                        cur_point_left = np.zeros(2)
                        strip_left = 0
                    else:
                        # print('----------- not smooth in left, slice: ', global_i, ', index: ', i, ', skip -------------')
                        pass
                except:
                    print('acos domain error')
                    print('last tan: ', last_tan_left)
                    print('cur tan: ', cur_tan_left)
                    print('dot: ', np.dot(last_tan_left, cur_tan_left))
                    print('mod * mod: ', np.linalg.norm(last_tan_left) * np.linalg.norm(cur_tan_left))
                    print('cos: ', np.dot(last_tan_left, cur_tan_left) / (np.linalg.norm(last_tan_left) * np.linalg.norm(cur_tan_left)))
                    exit(1)
            else:
                mean_tan_left += cur_tan_left
                mean_point_left += cur_point_left
                count_left += strip_left
                last_tan_left = cur_tan_left
                cur_tan_left = np.zeros(2)
                cur_point_left = np.zeros(2)
                strip_left = 0

        if i == round(overall_range_index[1]) - 1 or strip_right == strip_max:
            if (last_tan_right != np.zeros(2)).any():
                if np.dot(last_tan_right, cur_tan_right) / (np.linalg.norm(last_tan_right) * np.linalg.norm(cur_tan_right)) \
                        > math.cos(delta_theta):
                    mean_tan_right += cur_tan_right
                    mean_point_right += cur_point_right
                    count_right += strip_right
                    last_tan_right = cur_tan_right
                    cur_tan_right = np.zeros(2)
                    cur_point_right = np.zeros(2)
                    strip_right = 0
                else:
                    # print('----------- not smooth in right, slice: ', global_i, ', index: ', i,  ', skip -------------')
                    pass
            else:
                mean_tan_right += cur_tan_right
                mean_point_right += cur_point_right
                count_right += strip_right
                last_tan_right = cur_tan_right
                cur_tan_right = np.zeros(2)
                cur_point_right = np.zeros(2)
                strip_right = 0

    mean_tan_left /= count_left
    mean_tan_left /= np.linalg.norm(mean_tan_left)
    mean_point_left /= count_left

    mean_tan_right /= count_right
    mean_tan_right /= np.linalg.norm(mean_tan_right)
    mean_point_right /= count_right

    if ax is not None:
        draw_tangent(ax, mean_point_left, mean_tan_left, 1200)
        draw_tangent(ax, mean_point_right, mean_tan_right, 1200)

    return [[mean_point_left, mean_tan_left], [mean_point_right, mean_tan_right]]


def get_edges(image):
    """
    find the edges in the image by canny operator, return them.
    notes: always consider the input image as a CV-style coordinate system with (y, x) order

    :param image: label image to get edges from
    :return: edge_vertices, left_edges, right_edges
    """
    edges = feature.canny(image, sigma=2.0)
    # 所有边界像素点
    edge_vertices = []
    # 由下到上的两侧外边界点
    outer_left = []
    outer_right = []

    for i in range(image.shape[0] - 1, -1, -1):
        outer_found = False
        right = None
        for j in range(image.shape[1]):
            if edges[i, j]:
                edge_vertices.append([i, j])
                right = [i, j]
                if not outer_found:
                    outer_left.append([i, j])
                    outer_found = True
        if right is not None:
            outer_right.append(right)
    return edge_vertices, outer_left, outer_right


def has_placket(outer_left, outer_right, delta_x=30):
    for i in range(1, len(outer_left)):
        if abs(outer_left[i][1] - outer_left[i - 1][1]) > delta_x:
            return True
    for i in range(1, len(outer_right)):
        if abs(outer_right[i][1] - outer_right[i - 1][1]) > delta_x:
            return True
    if abs(outer_left[0][1] - outer_right[0][1]) > 20:
        return True
    return False


def middle_adjust(tan_image, edge_vertices, delta_x=30):
    """
    return the middle direction to adjust the long_axis_line

    :param tan_image:
    :param edge_vertices:
    :param delta_x: x threshold to judge the placket
    :return:
    """
    epical = get_epical(tan_image, edge_vertices)
    epical = np.mean(epical, axis=0)
    # epical.sort(key=lambda x: x[1])
    # epical = np.asarray(epical[0])

    edge_vertices = edge_vertices.copy()
    edge_vertices.sort(key=lambda x: x[0], reverse=True)
    bottom = []
    last_x = 0
    for x in edge_vertices:
        if len(bottom) == 0:
            bottom.append(np.asarray(x))
            last_x = x[1]
            continue
        if abs(x[1] - last_x) > delta_x:
            bottom.append(np.asarray(x))
            break
        last_x = x[1]

    bottom.sort(key=lambda x: x[1])
    # print('bottom are:', bottom)

    # correction for the inclination
    a, b = 1, 1
    bottom = (a * np.asarray(bottom[0]) + b * np.asarray(bottom[1])) / (a + b)
    dir = bottom - epical
    dir /= np.linalg.norm(dir)
    return dir


def get_long_axis_line(image, ratio_start=0.25, ratio_end=0.75) -> tuple:
    """
    find the long axis line in the given YX image

    :param image: always consider the input image as (Y, X) order
    :param ratio_start: same in get_overall_tangent
    :param ratio_end: same in get_overall_tangent
    :return: point, direction. both (2,) ndarray in (Y, X) order
    """
    tangent_image = get_tangent(image)
    edge_vertices, outer_left, outer_right = get_edges(image)
    left_line, right_line = get_overall_tangent(tangent_image, outer_left, outer_right, ratio_start, ratio_end)
    # 得到两侧切线的交点
    intersect_x = geoutils.lineLineIntersect(right_line[0], right_line[1], left_line[0], left_line[1])
    if (intersect_x - left_line[0]).dot(left_line[1]) > 0:
        left_line[1] *= -1
    if (intersect_x - right_line[0]).dot(right_line[1]) > 0:
        right_line[1] *= -1
    # 得到切线角分线的向下方向向量，作为长轴的方向向量
    axis_dir = left_line[1] + right_line[1]
    axis_dir /= np.linalg.norm(axis_dir)

    if has_placket(outer_left, outer_right):
        # TODO: may get the wrong middle_dir?
        middle_dir = middle_adjust(tangent_image, edge_vertices)
        l1, l2 = 3, 2
        axis_dir = (l1 * axis_dir + l2 * middle_dir) / (l1 + l2)
        axis_dir /= np.linalg.norm(axis_dir)
        pass

    return intersect_x, axis_dir


def get_exist_range(image_seq: np.ndarray) -> tuple:
    """
    find the existing interval of the labels, [start, end)

    :param image_seq: the bool-image sequences according with order (z, , )
    :return: the interval of z
    :raise:
    """
    start, end = 0, 0
    for i in range(image_seq.shape[0]):
        if np.sum(image_seq[i]) != 0:
            start = i
            break
    for i in range(image_seq.shape[0] - 1, -1, -1):
        if np.sum(image_seq[i] != 0):
            end = i + 1
            break
    if start == end:
        print("something wrong happened in get_exist_range, got start == end == ", start)
    return start, end


def get_long_axis_plane(image_seq: np.ndarray, ratio_down=0.3, ratio_up=0.6, ratio_start=0.25, ratio_end=0.75):
    """
    calculate the long axis plane by blending [ratio_down ~ ratio_up) YX images selected by their z value.

    :param image_seq: 3D image with the order (Z, Y, X), treated as 2D YX-image sequence
    :param ratio_down: lower bound of calculating slices
    :param ratio_up: upper bound of calculating slices
    :return:  p(point on the plane), nv(normal vector). the plane is expressed by XP dot nv = 0. notes: p and nv are (X, Y, Z) order
    """
    global global_i
    start, end = get_exist_range(image_seq)
    length = end - start
    temp = start

    start = temp + round(ratio_down * length)
    end = temp + round(ratio_up * length)

    print('start: ', start, ', end: ', end)

    blended_point = np.zeros(2)
    blended_dir = np.zeros(2)
    count = 0
    for i in range(start, end):
        global_i = i
        point, dir = get_long_axis_line(image_seq[i], ratio_start, ratio_end)
        if np.isnan(point).any() or np.isnan(dir).any():
            print('+++ skip nan +++')
            continue
        blended_point += point
        blended_dir += dir
        count += 1
    blended_point /= count
    blended_dir /= count
    blended_dir /= np.linalg.norm(blended_dir)
    if blended_point[0] < 0:
        blended_point += -blended_point[0] / blended_dir[0] * blended_dir
    return np.asarray([blended_point[1], blended_point[0], 0]), np.asarray([-blended_dir[0], blended_dir[1], 0])


def get_long_axis(label_array: np.ndarray, ratio_down=0.3, ratio_up=0.6, ratio_start=0.25, ratio_end=0.75) -> tuple:
    """
    get 3D long axis line by calculating the intersection of axis planes from transverse and sagittal plane

    :param label_array: origin 3D label input array, in order of (z, y, x)
    :param ratio_down: lower bound of calculating slices
    :param ratio_up: upper bound of calculating slices
    :param ratio_start: start Y index in each YX image slice
    :param ratio_end: end Y index in each YX image slice
    :return: point, direction in XYZ order
    """
    # (p0, n0): axis plane from transverse plane
    # (p1, n1): axis plane from sagittal plane

    p0, n0 = get_long_axis_plane(label_array, ratio_down, ratio_up, ratio_start, ratio_end)
    print('XY slices finished')


    label_array = label_array.swapaxes(0, 2)    # order is (X, Y, Z)
    label_array = np.flip(label_array, 2)   # order is (X, Y, Zmax - Z)
    p1, n1 = get_long_axis_plane(label_array, ratio_down, ratio_up, ratio_start, ratio_end)

    # change p1, n1 to be in the world 3D (X, Y, Z) coordinate system
    p1[2] = label_array.shape[2] - p1[0]
    p1[0] = 0
    n1[2] = -n1[0]
    n1[0] = 0

    axis_point, axis_direction = geoutils.planePlaneIntersect(p0, n0, p1, n1)
    return axis_point, axis_direction


def test():
    # 测试目录的本地路径
    # label_path = "..\\after_post_process1"
    # label_path = os.path.join(label_path, os.listdir(label_path)[0])
    label_path = r'C:\study\FFR\myProject\mytest\mydata_test\input\new_results_anti\prediction_00177.nii.gz'

    label_image = sitk.ReadImage(label_path)
    label_array = sitk.GetArrayFromImage(label_image)
    label_array = label_array.astype(np.float32)

    # TODO: test the sagittal plane
    # label_array = label_array.swapaxes(0, 2)
    # label_array = np.flip(label_array, 2)

    test_slice = 170   # transverse: 160; sagittal: 260

    # tangent_image = get_tangent(label_array[test_slice])
    # edge_vertices, outer_left, outer_right = get_edges(label_array[test_slice])

    fig, ax = plt.subplots()

    # TODO: test rotation
    label_array = label_array.swapaxes(1, 2)
    center, angle = get_rotation(label_array, 4 / 9, 5 / 9)
    print(center, angle)
    trans = sitk.Euler3DTransform()
    trans.SetCenter(np.asarray(center) * np.asarray(label_image.GetSpacing()) + np.asarray(label_image.GetOrigin()))
    trans.SetRotation(0, 0, np.deg2rad(angle))
    resampled = sitk.Resample(label_image, trans)

    label_array = sitk.GetArrayFromImage(resampled)
    label_array = label_array.swapaxes(0, 2)
    center, angle = get_rotation(label_array, 4 / 9, 5 / 9)
    trans.SetCenter(np.asarray((center[2], center[0], center[1])) * np.asarray(label_image.GetSpacing()) + np.asarray(label_image.GetOrigin()))
    trans.SetRotation(np.deg2rad(angle), 0, 0)
    resampled = sitk.Resample(resampled, trans)
    resampled.CopyInformation(label_image)
    #
    # sitk.WriteImage(resampled, 'test_resampled.nii.gz')
    label_array = sitk.GetArrayFromImage(resampled)
    label_array = label_array.astype(np.float32)
    test_slice = 100


    # TODO: test the edge vertices
    # print('left len: ', len(outer_left))
    # print('right len: ', len(outer_right))
    # ax.scatter([x[1] for x in outer_right], [x[0] for x in outer_right], s=[i for i in range(len(outer_right))], c=[[0, 1, 0]], marker='.')

    # TODO: test to mark the apical
    # epical = get_epical(tangent_image, edge_vertices, ax)

    # TODO: test to calculate the overall tangent
    # left_line, right_line = get_overall_tangent(tangent_image, outer_left, outer_right, ratio_start=0.4, ax=ax)
    # print('left line:', left_line)
    # print('right line:', right_line)

    # TODO: test the long axis of one image
    x, d = get_long_axis_line(label_array[test_slice])
    print('point: ', x)
    print('direction: ', d)
    # draw_tangent(ax, x, d, 300)

    # TODO: test the long axis plane
    # t1 = time.process_time()
    # p, n = get_long_axis_plane(label_array)
    # t2 = time.process_time()
    # print('one traverse of long axis plane costs time(s): ', t2 - t1)
    # point = np.asarray([p[1], p[0]])
    # direction = np.asarray([-n[0], n[1]])
    # print('point: ', point)
    # print('direction: ', direction)
    # draw_tangent(ax, point, direction, 300)

    # TODO: test the long axis line
    # t1 = time.process_time()
    # axis_point, axis_direction = get_long_axis(label_array, 1/3, 2/3, 0.4, 0.75)
    # t2 = time.process_time()
    # print('axis point: ', axis_point)
    # print('axis direction: ', axis_direction)
    # print('time cost(s): ', t2 - t1)

    # cached value for case_00177
    # axis_point = np.asarray([345.46317048, 9.9502225, 115.21751065])
    # axis_direction = np.asarray([-0.21520773, 0.97033094, 0.11019755])

    # TODO: test label with axis output. need long axis line test opened.
    # WARNING: maybe not work because of the uncertainty in point and direction
    # step = 2
    # temp_point = axis_point.copy()
    # for i in range(200):
    #     ind_i = round(temp_point[2])
    #     ind_j = round(temp_point[1])
    #     ind_k = round(temp_point[0])
    #     for j in range(-2, 3):
    #         for k in range(-2, 3):
    #             label_array[ind_i, ind_j + j, ind_k + k] = 2.0
    #     temp_point += step * axis_direction
    #
    # out_image = sitk.GetImageFromArray(label_array)
    # out_image.CopyInformation(label_image)
    # out_path = 'out_label_' + label_number + '.nii.gz'
    # sitk.WriteImage(out_image, out_path)


    # TODO: test the middle adjust
    tangent_image = get_tangent(label_array[test_slice])
    edge_vertices, outer_left, outer_right = get_edges(label_array[test_slice])
    epical = get_epical(tangent_image, edge_vertices, ax=ax)

    print('has placket:', has_placket(outer_left, outer_right))
    edge_input = outer_left + outer_right
    middle = middle_adjust(tangent_image, edge_input)
    epical = np.mean(epical, axis=0)
    draw_tangent(ax, epical, middle, half_length=200)

    ax.imshow(label_array[test_slice], cmap='Reds')
    plt.show()


def get_rotation(image_array: np.ndarray, ratio_down, ratio_up) -> tuple:
    """
    calculate the rotation angle and center from the XY slices

    :param image_array: 3D array seen as ZXY
    :param ratio_down: lower bound along the Z axis
    :param ratio_up:
    :return: point(XYZ), angle(along z axis)
    """
    start, end = get_exist_range(image_array)
    length = end - start
    temp = start
    start = temp + ratio_down * length
    end = temp + ratio_up * length

    center = [0., 0.]
    angle = 0.
    cnt = 0
    for i in range(round(start), round(end)):
        slice = image_array[i]
        slice = slice.astype(np.uint8)
        slice = slice[:, :, np.newaxis]
        contours, _ = cv2.findContours(slice, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        rect = cv2.minAreaRect(contours[0])
        center[0] += rect[0][0]
        center[1] += rect[0][1]
        angle += 90 - rect[-1]  # TODO: may have something wrong?
        cnt += 1

    center[0] /= cnt
    center[1] /= cnt
    angle /= cnt
    return (center[0], center[1], (start + end) / 2), angle


def execute_rotation(label_image, ratio_down, ratio_up):
    """
    execute rotation to create a new image array to be used in axis searching.

    :param label_image: original input image
    :param ratio_down:
    :param ratio_up:
    :return: rotated_array (ZYX order), theta_z (radian), center_z, theta_x, center_x as rotation order
    """
    label_array = sitk.GetArrayFromImage(label_image)
    label_array = label_array.astype(np.float32)
    label_array = label_array.swapaxes(1, 2)
    center, angle = get_rotation(label_array, ratio_down, ratio_up)
    theta_z = np.deg2rad(angle)
    center_z = center
    trans = sitk.Euler3DTransform()
    trans.SetCenter(np.asarray(center_z) * np.asarray(label_image.GetSpacing()) + np.asarray(label_image.GetOrigin()))
    trans.SetRotation(0, 0, theta_z)
    resampled = sitk.Resample(label_image, trans)

    label_array = sitk.GetArrayFromImage(resampled)
    label_array = label_array.astype(np.float32)

    if np.any(label_array) == 0:
        print('empty array!')

    label_array = label_array.swapaxes(0, 2)
    center, angle = get_rotation(label_array, ratio_down, ratio_up)
    theta_x = np.deg2rad(angle)
    center_x = (center[2], center[0], center[1])
    trans.SetCenter(np.asarray(center_x) * np.asarray(label_image.GetSpacing()) + np.asarray(label_image.GetOrigin()))
    trans.SetRotation(theta_x, 0, 0)
    resampled = sitk.Resample(resampled, trans)
    rotated_array = sitk.GetArrayFromImage(resampled)
    rotated_array = rotated_array.astype(np.float32)
    return rotated_array, theta_z, center_z, theta_x, center_x


def find_long_axis(image_path, ratio_down=1/3, ratio_up=2/3, ratio_start=0.25, ratio_end=0.75, out_path=None, out_rotated_path=None):
    """
    get 3D long axis line by calculating the intersection of axis planes from transverse and sagittal plane

    :param label_array: origin 3D label input array, in order of (z, y, x)
    :param ratio_down: lower bound of calculating slices
    :param ratio_up: upper bound of calculating slices
    :param ratio_start: start Y index in each YX image slice
    :param ratio_end: end Y index in each YX image slice
    :param out_path: path to write the image with long axis back. if None, no writing
    :param out_rotated_path: path to write the rotated image back.
    :return: point, direction in XYZ order
    """

    label_image = sitk.ReadImage(image_path)
    # label_array = sitk.GetArrayFromImage(label_image)
    # label_array = label_array.astype(np.float32)

    # TODO: handle pre_rotation and rotated image writing
    label_array, theta_z, center_z, theta_x, center_x = execute_rotation(label_image, ratio_down, ratio_up)


    t1 = time.process_time()
    axis_point, axis_direction = get_long_axis(label_array, ratio_down, ratio_up, ratio_start, ratio_end)
    t2 = time.process_time()

    print('axis point: ', axis_point)
    print('axis direction: ', axis_direction)
    print('time cost(s): ', t2 - t1)

    print('rotation history: ')
    print('theta_z', theta_z)
    print('center_z', center_z)
    print('theta_x', theta_x)
    print('center_x', center_x)

    if out_rotated_path is not None:
        step = 2
        temp_point = axis_point.copy()
        for i in range(200):
            ind_i = round(temp_point[2])
            ind_j = round(temp_point[1])
            ind_k = round(temp_point[0])
            for j in range(-2, 3):
                for k in range(-2, 3):
                    try:
                        label_array[ind_i, ind_j + j, ind_k + k] = 2.0
                    except:
                        continue
            temp_point += step * axis_direction

        out_image = sitk.GetImageFromArray(label_array)
        out_image.CopyInformation(label_image)
        sitk.WriteImage(out_image, out_rotated_path)

    # rotate the array and axis back
    # corrected..?
    label_array = sitk.GetArrayFromImage(label_image)
    label_array = label_array.astype(np.float32)
    axis_point -= np.asarray(list(center_x))
    R = np.asarray([
        [1, 0, 0],
        [0, np.cos(theta_x), -np.sin(theta_x)],
        [0, np.sin(theta_x), np.cos(theta_x)]
    ])
    axis_point = np.matmul(R, axis_point)
    axis_direction = np.matmul(R, axis_direction)
    axis_point += np.asarray(list(center_x))

    axis_point -= np.asarray(list(center_z))
    R = np.asarray([
        [np.cos(theta_z), -np.sin(theta_z), 0],
        [np.sin(theta_z), np.cos(theta_z), 0],
        [0, 0, 1]
    ])
    axis_point = np.matmul(R, axis_point)
    axis_direction = np.matmul(R, axis_direction)
    axis_point += np.asarray(list(center_z))

    print('rotated back')
    print('axis point: ', axis_point)
    print('axis direction: ', axis_direction)

    if out_path is not None:
        # use the function to draw axis in the array
        label_array = draw_axis(label_array, axis_point, axis_direction, 2.0)
        out_image = sitk.GetImageFromArray(label_array)
        out_image.CopyInformation(label_image)
        sitk.WriteImage(out_image, out_path)

    return axis_point, axis_direction


def draw_axis(label_array, axis_point, axis_direction, tag):
    """
    draw the long axis into the label array with the value $tag
    :param label_array:
    :param axis_point:
    :param axis_direction:
    :return:
    """
    step = 2
    temp_point = axis_point.copy()
    for i in range(200):
        ind_i = round(temp_point[2])
        ind_j = round(temp_point[1])
        ind_k = round(temp_point[0])
        for j in range(-2, 3):
            for k in range(-2, 3):
                try:
                    label_array[ind_i, ind_j + j, ind_k + k] = tag
                except:
                    continue
        temp_point += step * axis_direction
    return label_array


if __name__ == '__main__':
    np.set_printoptions(precision=16)
    test()

    # label_path = 'C:\\study\\FFR\\myProject\\mytest\\mydata\\input\\prediction_00177.nii.gz'
    # out_path = 'out_all.nii.gz'
    # out_ro_pa = 'out_all_rotationed.nii.gz'
    # ratio_down = 4/9
    # ratio_up = 5/9
    # ratio_start = 0.4  # 0.5?
    # ratio_end = 0.75
    #
    # point, vector = find_long_axis(label_path, ratio_down, ratio_up, ratio_start, ratio_end, out_path=out_path, out_rotated_path=out_ro_pa)
